#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
模型训练脚本
从数据库获取训练数据并微调BERT模型
"""

import sys
import os
import json
import logging
import torch
from datetime import datetime
from database_manager import DatabaseManager
from bert_finetuner import BERTFineTuner

# 配置日志
# 确保日志目录存在
log_dir = "logs"
if not os.path.exists(log_dir):
    os.makedirs(log_dir)

# 为这个模块创建独立的logger
logger = logging.getLogger('train_model')
logger.setLevel(logging.DEBUG)

# 清除可能存在的旧处理器
for handler in logger.handlers[:]:
    logger.removeHandler(handler)

# 创建文件处理器
file_handler = logging.FileHandler(os.path.join(log_dir, "training.log"), encoding='utf-8')
file_handler.setLevel(logging.DEBUG)

# 创建控制台处理器
console_handler = logging.StreamHandler(sys.stdout)
console_handler.setLevel(logging.INFO)

# 创建格式化器
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
file_handler.setFormatter(formatter)
console_handler.setFormatter(formatter)

# 添加处理器到logger
logger.addHandler(file_handler)
logger.addHandler(console_handler)

# 防止日志传播到根logger
logger.propagate = False

# 获取当前模块的logger
current_logger = logging.getLogger('train_model')

def train_model():
    """
    训练模型主函数
    """
    current_logger.info("=" * 50)
    current_logger.info("开始模型训练流程...")
    current_logger.info("=" * 50)
    
    # 检查CUDA是否可用
    use_cuda = torch.cuda.is_available()
    device_count = torch.cuda.device_count()
    current_logger.info(f"CUDA可用: {use_cuda}, 设备数量: {device_count}")
    
    if use_cuda:
        for i in range(device_count):
            current_logger.info(f"  设备 {i}: {torch.cuda.get_device_name(i)}")
    else:
        current_logger.warning("CUDA不可用，将在CPU上进行训练")
    
    # 初始化数据库管理器
    db_manager = DatabaseManager()
    
    # 连接数据库
    current_logger.info("连接数据库...")
    if not db_manager.connect():
        current_logger.error("错误: 无法连接到数据库")
        return False
    
    # 获取训练数据
    current_logger.info("获取训练数据...")
    train_data = db_manager.get_bert_train_data()
    
    # 断开数据库连接
    db_manager.disconnect()
    
    if not train_data:
        current_logger.error("错误: 没有获取到训练数据")
        return False
    
    current_logger.info(f"成功获取 {len(train_data)} 条训练数据")
    
    # 显示数据结构信息
    current_logger.info("数据结构分析:")
    valid_records = 0
    field_ids = set()
    
    for i, data in enumerate(train_data[:10]):  # 分析前10条记录
        current_logger.info(f"  记录 {i+1}:")
        current_logger.info(f"    ID: {data.get('id', 'N/A')}")
        current_logger.info(f"    Field ID: {data.get('field_id', 'N/A')}")
        if data.get('field_id'):
            field_ids.add(data['field_id'])
        
        current_logger.info(f"    实体信息类型: {type(data.get('entity_info'))}")
        if data.get('entity_info'):
            if isinstance(data['entity_info'], dict):
                current_logger.info(f"    包含文本: {data['entity_info'].get('text', 'N/A')}")
                entities = data['entity_info'].get('entities', [])
                current_logger.info(f"    实体数量: {len(entities) if isinstance(entities, list) else 'N/A'}")
                if isinstance(entities, list):
                    for j, entity in enumerate(entities[:3]):  # 只显示前3个实体
                        current_logger.info(f"      实体 {j+1}: {entity}")
            else:
                current_logger.info(f"    实体信息内容: {data['entity_info']}")
        
        # 检查是否为有效记录
        entity_info = data.get('entity_info')
        if (entity_info and isinstance(entity_info, dict) and 
            'text' in entity_info and 'entities' in entity_info and
            isinstance(entity_info['entities'], list)):
            valid_records += 1
    
    current_logger.info(f"有效记录数量: {valid_records}/{min(10, len(train_data))} (在前10条记录中)")
    if field_ids:
        current_logger.info(f"涉及的领域ID: {field_ids}")
    
    # 如果有效记录太少，给出警告
    if valid_records == 0:
        current_logger.warning("前10条记录中没有有效的训练数据")
        response = input("是否继续训练? (y/N): ")
        if response.lower() != 'y':
            current_logger.info("取消训练")
            return True
    
    # 过滤有效数据
    filtered_train_data = []
    invalid_count = 0
    for data in train_data:
        entity_info = data.get('entity_info')
        if (entity_info and isinstance(entity_info, dict) and 
            'text' in entity_info and 'entities' in entity_info and
            isinstance(entity_info['entities'], list)):
            filtered_train_data.append(data)
        else:
            invalid_count += 1
            current_logger.debug(f"过滤掉无效记录: {data.get('id', 'N/A')}")
    
    if not filtered_train_data:
        current_logger.error("错误: 没有有效的训练数据用于微调")
        return False
    
    current_logger.info(f"过滤后有效数据数量: {len(filtered_train_data)} (过滤掉 {invalid_count} 条无效数据)")
    
    # 显示过滤后的数据样本
    current_logger.info("过滤后数据样本:")
    for i, data in enumerate(filtered_train_data[:3]):
        entity_info = data.get('entity_info', {})
        text = entity_info.get('text', '')
        entities = entity_info.get('entities', [])
        current_logger.info(f"  样本 {i+1}:")
        current_logger.info(f"    文本: {text}")
        current_logger.info(f"    实体: {entities}")
    
    # 初始化微调器
    current_logger.info("初始化BERT微调器...")
    try:
        finetuner = BERTFineTuner()
    except Exception as e:
        current_logger.error(f"错误: 初始化微调器失败: {e}")
        import traceback
        current_logger.error(traceback.format_exc())
        return False
    
    # 创建输出目录
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir = f"./fine_tuned_model_{timestamp}"
    
    # 开始微调
    current_logger.info(f"开始微调模型，输出目录: {output_dir}")
    try:
        # 设置环境变量以获得更详细的CUDA错误信息
        os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
        os.environ['TORCH_USE_CUDA_DSA'] = '1'
        
        model_path = finetuner.fine_tune(
            train_data=filtered_train_data,
            output_dir=output_dir,
            num_train_epochs=10,  # 增加训练轮数以提高效果
            per_device_train_batch_size=4,  # 调整批次大小以提高训练效果
            save_steps=50
        )
        current_logger.info("模型微调完成!")
        current_logger.info(f"微调后的模型保存在: {model_path}")
        return True
    except Exception as e:
        current_logger.error(f"错误: 模型微调失败: {e}")
        import traceback
        current_logger.error(traceback.format_exc())
        return False

def test_fine_tuned_model(model_path: str = "./fine_tuned_model"):
    """
    测试微调后的模型
    
    Args:
        model_path: 微调模型的路径
    """
    current_logger.info("测试微调后的模型...")
    
    try:
        from fine_tuned_ner import FineTunedNERProcessor
        
        # 检查模型路径是否存在
        if not os.path.exists(model_path):
            current_logger.error(f"错误: 模型路径不存在: {model_path}")
            return False
        
        # 加载微调模型
        processor = FineTunedNERProcessor(model_path)
        
        # 测试文本
        test_texts = [
            "承德钒钛钳工需要50MPa的压力容器操作证",
            "河钢集团实施人工智能+行动工作方案要求提高组织效率",
            "今天承德钒钛的钳工完成了350℃高温测试",
            # 添加一些训练数据中可能未出现但符合语言规律的实体
            "首钢的技术工程师正在研究新型材料",
            "宝钢集团的焊工需要具备相关资质证书",
            "测试温度需要达到600℃以上"
        ]
        
        current_logger.info("测试结果:")
        for text in test_texts:
            current_logger.info(f"文本: {text}")
            try:
                entities = processor.extract_entities(text)
                if entities:
                    for entity in entities:
                        current_logger.info(f"  {entity['nerKey']}: {entity['nerValue']} (置信度: {entity.get('confidence', 'N/A'):.3f})")
                else:
                    current_logger.info("  未识别到实体")
            except Exception as e:
                current_logger.error(f"  处理文本时出错: {e}")
                import traceback
                current_logger.error(traceback.format_exc())
        
        return True
    except Exception as e:
        current_logger.error(f"测试微调模型时出错: {e}")
        import traceback
        current_logger.error(traceback.format_exc())
        return False

def main():
    """
    主函数
    """
    if len(sys.argv) > 1:
        command = sys.argv[1]
        
        if command == "train":
            success = train_model()
            sys.exit(0 if success else 1)
        elif command == "test":
            model_path = sys.argv[2] if len(sys.argv) > 2 else "./fine_tuned_model"
            success = test_fine_tuned_model(model_path)
            sys.exit(0 if success else 1)
        else:
            print("未知命令。可用命令:")
            print("  train - 训练模型")
            print("  test [model_path] - 测试微调后的模型")
            sys.exit(1)
    else:
        print("用法:")
        print("  python train_model.py train - 训练模型")
        print("  python train_model.py test [model_path] - 测试微调后的模型")
        sys.exit(1)

if __name__ == "__main__":
    main()